package com.demo.java.kafka.consumer;

import lombok.extern.slf4j.Slf4j;
import org.apache.kafka.clients.consumer.*;
import org.apache.kafka.clients.producer.*;
import org.apache.kafka.common.TopicPartition;
import org.junit.Test;

import java.time.Duration;
import java.util.*;

@Slf4j
public class OffsetManageTest {

    @Test
    public void producerTest() {
        Properties prop = new Properties();
        prop.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092");
        prop.put(ProducerConfig.ENABLE_IDEMPOTENCE_CONFIG, "true");
        prop.put(ProducerConfig.TRANSACTIONAL_ID_CONFIG, "apptransmit-1");
        prop.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.StringSerializer");
        prop.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.StringSerializer");
        prop.put(ProducerConfig.RETRIES_CONFIG, "3");
        KafkaProducer<String, String> producer = new KafkaProducer<>(prop);

        producer.initTransactions();
        producer.beginTransaction();

        try {
            // do soamething
            String topic = "testoffset";
            Integer part = 0;
            for (int i = 0; i < 10; i++) {
                ProducerRecord<String, String> record = new ProducerRecord<>(topic, part, "key-" + i, "value-" + i);
                producer.send(record, new Callback() {
                    @Override
                    public void onCompletion(RecordMetadata recordMetadata, Exception e) {
                        System.out.println("success");
                        System.out.println(recordMetadata.topic() + ", " + recordMetadata.partition() + ", " + recordMetadata.offset());
                    }
                });
            }

            producer.commitTransaction();
        } catch (Exception ex) {
            log.error("==== transaction error, ", ex);
            producer.abortTransaction();
        } finally {
            if (producer != null) {
                producer.close();
            }
        }
    }

    @Test
    public void consumerAsyncCommitOffset() {
        Properties prop = new Properties();
        prop.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092");
        prop.put(ConsumerConfig.GROUP_ID_CONFIG, "apptransmit");
        prop.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false");
        prop.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");// latest, earliest, none
        prop.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.StringDeserializer");
        prop.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.StringDeserializer");
        KafkaConsumer<String, String> consumer = new KafkaConsumer(prop);
        consumer.subscribe(Collections.singleton("testoffset"));

        // 获取topic的分区信息
        Set<TopicPartition> assignment = consumer.assignment();
        printAssignment("subscribe", assignment);

        while (true) {
            ConsumerRecords<String, String> records = consumer.poll(Duration.ofSeconds(60));
            printAssignment("pool", consumer.assignment());
            Set<TopicPartition> partitions = records.partitions();
            for (TopicPartition partition : partitions) {
                final int part = partition.partition();
                final List<ConsumerRecord<String, String>> recordList = records.records(partition);
                for (ConsumerRecord<String, String> record : recordList) {
                    log.info("=== record offset: {}, partition: {}, topic: {}, key: {}, value: {}", record.offset(), part, record.topic(), record.key(), record.value());
                }

                ConsumerRecord<String, String> lastConsumerRecord = recordList.get(recordList.size()-1);
                long lastOffset = lastConsumerRecord.offset();
                final Optional<Integer> integer = lastConsumerRecord.leaderEpoch();
                log.info("=== leaderEpoch: {}, offset: {}", integer.get(), lastOffset);
                // 同步提交offset
                consumer.commitSync(Collections.singletonMap(partition, new OffsetAndMetadata(lastOffset + 1)));
            }
        }
    }

    public void printAssignment(String prefix, Set<TopicPartition> assignmentSet) {
        for (TopicPartition partition : assignmentSet) {
            log.info("=== {} , assignment: partition: {}, topic: {}, {}", prefix, partition.partition(), partition.topic(), partition.toString());
        }
    }
}
